# -*- coding: utf-8 -*-
"""
Created on Sat Aug 17 19:09:28 2024

@author: TAO

图片分析程序，自动寻找并标记畸变图上的参考点，根据模板图片旋转，平移，和缩放图片
矫正成正交图片，并在正交图片上寻找小球的坐标，并返回成 excel 文件

修改记录：
     2024.8.28  删除标识点坐标输出，excel文件中只包含文件名，小球的X,Y像素坐标
                删除excel文件中的首行模板记录
                把输出的坐标系从左上角0,0点转换为左下角0,0点，符合常规习惯
"""

import cv2
import numpy as np
import os
import glob
import pandas as pd

def approx_contour(contour):
    """
    将轮廓近似为一个多边形。
    参数 contour: 输入的轮廓，是由一系列点构成的列表或 NumPy 数组。
    返回:近似后的多边形轮廓。
    """
    epsilon = 0.02 * cv2.arcLength(contour, True)  # 计算轮廓的周长，并基于此确定近似精度
    return cv2.approxPolyDP(contour, epsilon, True)  # 使用近似算法获取多边形轮廓

def is_quadrilateral(cnt):
    """
    检查轮廓是否为四边形，用来搜索图片上的识别点   
    参数:cnt: 输入的轮廓。        
    返回:如果轮廓为四边形且为凸多边形，则返回 True；否则返回 False。
    """
    approx = approx_contour(cnt)  # 将轮廓近似为多边形
    return len(approx) == 4 and cv2.isContourConvex(approx)  # 检查近似后的多边形是否为四边形且凸

def is_ellipse(cnt):
    """
    检查轮廓是否可近似为椭圆，并返回椭圆的信息，这里是搜索图片上的小球    
    参数 cnt: 输入的轮廓。        
    返回:一个元组 (is_ellipse, ellipse_info)，其中 is_ellipse 是布尔值，指示轮廓是否为椭圆，
        ellipse_info 是字典，包含椭圆的详细信息。
    """
    area = cv2.contourArea(cnt)  # 计算轮廓的面积
    if area < 100:  # 如果轮廓面积太小，则不认为是椭圆
        ellipse_info = {
            "center": None,
            "axes": None,
            "angle": None,
            "eccentricity": None
        }
        return False, ellipse_info
    
    (x, y), (MA, ma), angle = cv2.fitEllipse(cnt)  # 使用椭圆拟合轮廓
    eccentricity = np.sqrt(1 - min(MA, ma) ** 2 / max(MA, ma) ** 2)  # 计算椭圆的偏心率
    
    # 检查是否为椭圆
    is_ellipse = eccentricity < 0.8
    
    # 返回字典
    ellipse_info = {
        "center": (x, y),
        "axes": (MA, ma),
        "angle": angle,
        "eccentricity": eccentricity
    }
    
    return is_ellipse, ellipse_info

# 排序质心以形成凸四边形
# 我们使用Jarvis March算法（也称为Gift Wrapping算法）来找到凸包
def gift_wrapping(points):
# 找到最左下角的点作为起始点
    start = min(points, key=lambda p: (p[1], p[0]))
    hull = [start]
    point = start
    
    # 主循环
    while True:
        # 初始化下一个点为None
        next_point = None
        # 遍历所有点，找到使ccw最大的点
        for candidate in points:
            # 跳过当前点
            if candidate == point:
                continue
            # 如果没有找到下一个点，或者当前候选点使ccw更大，则更新next_point
            if next_point is None or ccw(point, candidate, next_point) > 0:
                next_point = candidate
        
        # 如果没有找到下一个点，或者下一个点回到了起点，则退出循环
        if next_point is None or next_point == start:
            break
        
        # 添加下一个点到hull并更新point
        hull.append(next_point)
        point = next_point
    
    return hull

def ccw(A, B, C):
    return (C[1]-A[1]) * (B[0]-A[0]) > (B[1]-A[1]) * (C[0]-A[0])

#透视变换
def perspective_transform(src_image, src_points, dst_points):
    """
    使用OpenCV进行透视变换。
    :参数 src_image: 需要变形的源图像。
    :参数 src_points: 对应于模板图像角点的目标点列表。
    :参数 dst_points: 模板图像的角点列表。
    :返回: 变换后的图像。
    """

    # 计算透视变换矩阵
    transform_matrix = cv2.getPerspectiveTransform(np.array(src_points, dtype=np.float32), np.array(dst_points, dtype=np.float32))

    # 应用透视变换
    corrected_image = cv2.warpPerspective(src_image, transform_matrix, (src_image.shape[1], src_image.shape[0]))

    return corrected_image
#寻找图片上的4边形
def find_and_sort_corners(image, contours):
    corners = []
    areas = []
    first_ellipse_center = None  # 用于存储第一个椭圆的中心坐标

    markedimage = image.copy()
    for contour in contours:
        # 确定四边形
        #if len(approx) == 4:
        if is_quadrilateral(contour):
            M = cv2.moments(contour)
            cx = int(M['m10']/M['m00'])
            cy = int(M['m01']/M['m00'])
            area = cv2.contourArea(contour)
            corners.append((cx, cy))
            areas.append(area)
            # 在质心画直径5像素的实心点
            cv2.circle(markedimage, (cx, cy), 5, (0, 0, 255), -1)
            cv2.drawContours(markedimage, [contour], 0, (0, 255, 0), 2)
        else:
            # 如果不是四边形，则检查是否为椭圆
            b_ellipse, ellipse_info = is_ellipse(contour)
            if b_ellipse and first_ellipse_center is None:
                ecenter = ellipse_info.get('center')
                if ecenter is not None:
                    first_ellipse_center = (int(ecenter[0]), int(ecenter[1]))
                    cv2.drawContours(markedimage, [contour], 0, (0, 255, 0), 2)
                    cv2.circle(markedimage, first_ellipse_center, 5, (0, 0, 255), -1)

    # 找出最小的四边形D
    D_index = areas.index(min(areas))
    D_point = corners[D_index]
    
     # 获取凸包
    convex_hull = gift_wrapping(corners)
    # 找到D点在凸包中的位置
    D_hull_index = convex_hull.index(D_point)
    
    # 连接凸包点
    for i in range(len(convex_hull)):
        next_i = (i + 1) % len(convex_hull)
        cv2.line(markedimage, convex_hull[i], convex_hull[next_i], (0, 255, 0), 2)

    # 构建字典保存字母和实际坐标
    corners_dict = {'D': D_point}
    # 循环遍历凸包点，找到A、B、C点
    for corner_name, offset in zip(['A', 'B', 'C'], [1, 2, 3]):
        next_index = (D_hull_index + offset) % len(convex_hull)
        corners_dict[corner_name] = convex_hull[next_index]

    # 输出和标记质心
    for corner_name, corner in corners_dict.items():
        # 计算文本位置，使其在质心正下方
        text_position = (corner[0], corner[1] + 80)  # 假设我们想向下偏移80像素
        cv2.putText(markedimage, corner_name, text_position, cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
        print(f"Corner {corner_name}: {corner}")

    outPoints = [corners_dict['A'], corners_dict['B'], corners_dict['C'], corners_dict['D']]

    return image,markedimage,outPoints,first_ellipse_center

#查找指定图片上的多边形,并返回坐标
def get_coordinates_from_image(image_path):
    img = cv2.imread(image_path)

    # 转换为灰度图像
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 高斯模糊去噪
    blurred_img = cv2.GaussianBlur(gray_img, (5, 5), 0)

    # 二值化处理
    threshold_value = 60
    ret, binary_img = cv2.threshold(blurred_img, threshold_value, 255, cv2.THRESH_BINARY_INV)

    # 轮廓检测
    contours, _ = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 搜索关键坐标，并标记图片
    return find_and_sort_corners(img, contours)

def find_ellipse_center_point(img):
    first_ellipse_center = None
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 高斯模糊去噪
    blurred_img = cv2.GaussianBlur(gray_img, (5, 5), 0)
    # 二值化处理
    threshold_value = 60
    ret, binary_img = cv2.threshold(blurred_img, threshold_value, 255, cv2.THRESH_BINARY_INV)
    contours, _ = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        if is_quadrilateral(contour):
            M = cv2.moments(contour)
            #cx = int(M['m10']/M['m00'])
            #cy = int(M['m01']/M['m00'])
            # 在质心画直径5像素的实心点
            #cv2.circle(img, (cx, cy), 5, (0, 0, 255), -1)
            #cv2.drawContours(img, [contour], 0, (0, 255, 0), 2)
        else:
            b_ellipse, ellipse_info = is_ellipse(contour)
            if b_ellipse and first_ellipse_center is None:
                ecenter = ellipse_info.get('center')
                if ecenter is not None:
                    first_ellipse_center = (int(ecenter[0]), int(ecenter[1]))
                    cv2.drawContours(img, [contour], 0, (0, 255, 0), 2)
                    cv2.circle(img, first_ellipse_center, 5, (0, 0, 255), -1)
                break
        
    return first_ellipse_center   
#------------------------------------------------------------------------------
# 输出坐标到Excel 文件中
# Excel 的第一行数据是模板图片的坐标,前四组坐标分别对应参考标志的A,B,C,D点,第5组数据忽略
# Excel 的第2行到第N+2行是解析了N张图片的结果坐标,分别是畸变图的A,B,C,D点坐标,第5组数据是小球的坐标
# template_image_file 是模板图片，标准的正交图片
# directory_path 需要遍历处理的保存畸变图片的目录路径
# output_excel_file 结果输出的Excel 文件名路径
def write_coordinates_to_excel(template_image_file,directory_path, output_excel_file):
   
    # 初始化DataFrame
    # 定义列名
    #columns = ['ImageName', 'AX', 'AY', 'BX', 'BY', 'CX', 'CY', 'DX', 'DY','BallX','BallY']
    # 删除标识点坐标
    columns = ['ImageName', 'BallX','BallY']
    df = pd.DataFrame(columns=columns)


    # 获取模板图片的A,B,C,D坐标,函数返回4个坐标，最后一个小球坐标忽略
    orgimage,markedImage,template_coordinates,ellipsepoint = get_coordinates_from_image(template_image_file)

    file_name = os.path.basename(template_image_file)
    
    file_name = f"T_{file_name}"
    
    DEBUG_MODE = False
    #DEBUG_BEGIN  输出标记后的模板文件, 注释DEBUG_MODE = True将关闭调试标记图片输出
    DEBUG_MODE = True
    # 添加 "_1" 到路径后面,输出标记后或者矫正后的图片
    corrected_path = f"{directory_path}_1"
    if DEBUG_MODE:
        # 创建矫正后的路径的目录（如果不存在）
        os.makedirs(corrected_path, exist_ok=True)    
        jpg_files = glob.glob(os.path.join(corrected_path, "*.jpg"))
    
        # 删除所有 .jpg 文件，上一次生成的文件
        for file_del_path in jpg_files:
            os.remove(file_del_path)
         
        corrected_file_path = os.path.join(corrected_path, file_name)
        cv2.imwrite(corrected_file_path, markedImage)
    #DEBUG_END      
    
    # 将坐标转换为DataFrame的一行
    #row_data = [file_name] + [coord for point in template_coordinates for coord in point] + list(ellipsepoint)
    #只输出小球的X,Y坐标，不再输出标识点坐标
    #row_data = [file_name] + list(ellipsepoint)
  
    #new_row = pd.Series(row_data, index=columns)

    # 使用 concat 添加行到 DataFrame
    #df = pd.concat([df, new_row.to_frame().T], ignore_index=True)
    #-----模板读取完毕-------    
    
    
    # 接下来遍历指定目录中的所有图片文件，并把结果添加到df中
    # 获取目录中所有文件名
    files = os.listdir(directory_path)
    
    # 创建一个字典来存储文件名及其对应的数字
    image_files = {}
    # 创建一个列表来存储文件名
    image_list = []
    for file in files:
        if file.startswith('Image') and file.endswith('.jpg'):
            try:
                number = int(file[5:-4])  # 提取文件名中的数字部分
                image_files[number] = file
                image_list.append(number)
            except ValueError:
                # 忽略不符合格式的文件
                pass
    
    # 按照数字排序,考虑到以数字命名的图片文件，数字可能不是连续的
    sorted_numbers = sorted(image_list)
    

        
    # 遍历文件
    for number in sorted_numbers:
        file_name = image_files[number]
        file_path = os.path.join(directory_path, file_name)
        
        
        # 获取坐标
        orgimage,markedImage,coordinates,ellipsepoint = get_coordinates_from_image(file_path)
        
        # 矫正图像,最关键的代码，透视变换
        corrected_image = perspective_transform(orgimage, coordinates,template_coordinates)

        # 从矫正后的图片中搜索小球坐标
        ellipsepoint = find_ellipse_center_point(corrected_image)
        
        # 获取图像的高度
        height, _ = corrected_image.shape[:2]
        # 把坐标系从左上角0,0点转换为左下角0,0点，由于X方向是一致的，
        # 因此只需要用图片的高度减去Y坐标，得到新的Y坐标
        ball_center = list(ellipsepoint)
        ball_center[1] = height - ball_center[1]
        
        #DEBUG_BEGIN  输出标记后的畸变矫正图片文件
        # 保存矫正后的图片
        if DEBUG_MODE:
            marked_file_name = f"M_{file_name}"
            marked_file_path = os.path.join(corrected_path, marked_file_name)
            corrected_file_path = os.path.join(corrected_path, file_name)
            cv2.imwrite(marked_file_path, markedImage)
            cv2.imwrite(corrected_file_path, corrected_image)
        #DEBUG_END     
        
        # 将坐标转换为DataFrame的一行
        #row_data = [file_name] + [coord for point in template_coordinates for coord in point]  + list(ellipsepoint)
        row_data = [file_name] + ball_center
        new_row = pd.Series(row_data, index=columns)

        # 使用 concat 添加行到 DataFrame
        df = pd.concat([df, new_row.to_frame().T], ignore_index=True)
    
    # 写入Excel文件
    df.to_excel(output_excel_file, index=False)
    print(f"数据已写入到 {output_excel_file}")

#------------------------------------------------------------------------------
# 主函数从这里开始
if __name__ == "__main__":
    #template_image_path 模板图片路径，获得标准图片的4个点
    template_image_path = 'E:/oball/Bin/20240908_141152/image1.jpg'
    #需要处理的畸变的图片路径，包含若干张图片
    input_image_directory= 'E:/oball/Bin/20240908_141152'
    # 指定输出的excel文件,保存在图片目录下
    output_excel_file = os.path.join(input_image_directory, 'coordinates.xlsx')
    
    # 调用函数遍历目标文件目录，产生坐标文件
    write_coordinates_to_excel(template_image_path,input_image_directory, output_excel_file)
    
